package com.pgp.demo.pgpexample.util;
import com.pgp.demo.pgpexample.exception.BusinessException;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.exception.ExceptionUtils;
import org.bouncycastle.bcpg.ArmoredOutputStream;
import org.bouncycastle.bcpg.BCPGOutputStream;
import org.bouncycastle.bcpg.HashAlgorithmTags;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.openpgp.*;
import org.bouncycastle.openpgp.operator.PBESecretKeyDecryptor;
import org.bouncycastle.openpgp.operator.PublicKeyDataDecryptorFactory;
import org.bouncycastle.openpgp.operator.jcajce.*;
import java.io.*;
import java.security.GeneralSecurityException;
import java.security.SecureRandom;
import java.security.Security;
import java.util.Iterator;
/**
 * pgp helper
 *
 * @author: alan
 * @date: 2024/12/03
 */
@Slf4j
public class PGPHelper {
    private static final String PROVIDER_NAME = "BC";

    static {
        Security.addProvider(new BouncyCastleProvider());
    }
    /**
     * 解密文件
     *
     * @param encryptFile  加密文件.gpg
     * @param originalFile 原始文件.csv
     * @param secretKey    私钥.asc
     * @param passphrase   密码 string
     * @throw {@link BusinessException}   文件解密异常
     */
    public static void decryptFile(InputStream encryptFile, OutputStream originalFile, InputStream secretKey, String passphrase) {
        if (encryptFile == null || originalFile == null || secretKey == null || passphrase == null || passphrase.isEmpty()) {
            throw new BusinessException("[decryptFile] 参数存在空值");
        }
        try {
            encryptFile = PGPUtil.getDecoderStream(encryptFile);
            PGPObjectFactory pgpF = new PGPObjectFactory(encryptFile);
            PGPEncryptedDataList enc;
            Object o = pgpF.nextObject();
            if (o instanceof  PGPEncryptedDataList) {
                enc = (PGPEncryptedDataList) o;
            } else {
                enc = (PGPEncryptedDataList) pgpF.nextObject();
            }
            Iterator it = enc.getEncryptedDataObjects();
            PGPPrivateKey sKey = null;
            PGPPublicKeyEncryptedData pbe = null;

            while (sKey == null && it.hasNext()) {
                pbe = (PGPPublicKeyEncryptedData) it.next();
                sKey = findSecretKey(secretKey, pbe.getKeyID(), passphrase.toCharArray());
            }

            if (sKey == null) {
                throw new IllegalArgumentException("[decryptFile] 未找到消息的密钥");
            }

            PublicKeyDataDecryptorFactory b = new JcePublicKeyDataDecryptorFactoryBuilder()
                    .setProvider(PROVIDER_NAME)
                    .setContentProvider(PROVIDER_NAME).build(sKey);

            InputStream clear = pbe.getDataStream(b);
            PGPObjectFactory plainFact = new PGPObjectFactory(clear);
            Object message = plainFact.nextObject();

            if (message instanceof  PGPCompressedData) {
                PGPCompressedData compressedData = (PGPCompressedData) message;
                PGPObjectFactory pgpFact = new PGPObjectFactory(compressedData.getDataStream());
                message = pgpFact.nextObject();
            }

            if (message instanceof  PGPLiteralData) {
                PGPLiteralData literalData = (PGPLiteralData) message;
                InputStream unc = literalData.getInputStream();
                int ch;
                while ((ch = unc.read()) >= 0) {
                    originalFile.write(ch);
                }
            } else if (message instanceof  PGPOnePassSignatureList) {
                throw new PGPException("[decryptFile] 加密消息包含签名消息，而不是文字数据。");
            } else {
                throw new PGPException("[decryptFile] 消息不是简单的加密文件，类型未知。");
            }

            if (pbe.isIntegrityProtected()) {
                if (!pbe.verify()) {
                    throw new PGPException("[decryptFile] 消息完整性检查失败");
                }
            }
        } catch (Exception e) {
            log.error("[decryptFile] error:", e);
            throw new BusinessException("500", "解密异常", ExceptionUtils.getStackTrace(e));
        } finally {
            try {
                encryptFile.close();
                originalFile.close();
                secretKey.close();
            } catch (Exception e) {
                log.error("[decryptFile.close] error:", e);
            }
        }
    }

    /**
     * 加密文件
     *
     * @param encryptFile   加密文件.gpg
     * @param originalFile  原始文件.csv
     * @param publicKey     公钥.asc
     * @param armor         装甲 可选(默认true)
     * @param withIntegrityCheck 使用完整性检查(默认true)
     * @throw {@link BusinessException}     文件加密异常
     */
    public static void encryptFile(OutputStream encryptFile, File originalFile, InputStream publicKey, boolean armor, boolean withIntegrityCheck) {
        if (encryptFile == null || originalFile == null || publicKey == null) {
            throw new BusinessException("[encryptFile] 参数存在空值");
        }
        OutputStream cOut = null;
        try {
            PGPPublicKey pgpPublicKey = readPublicKey(publicKey);
            if (armor) {
                encryptFile = new ArmoredOutputStream(encryptFile);
            }
            ByteArrayOutputStream bOut = new ByteArrayOutputStream();
            PGPCompressedDataGenerator comData = new PGPCompressedDataGenerator(PGPCompressedData.ZIP);
            PGPUtil.writeFileToLiteralData(comData.open(bOut), PGPLiteralData.BINARY, originalFile);
            comData.close();
            JcePGPDataEncryptorBuilder c = new JcePGPDataEncryptorBuilder(PGPEncryptedData.CAST5).setWithIntegrityPacket(withIntegrityCheck)
                    .setSecureRandom(new SecureRandom())
                    .setProvider(PROVIDER_NAME);

            PGPEncryptedDataGenerator cPk = new PGPEncryptedDataGenerator(c);
            JcePublicKeyKeyEncryptionMethodGenerator d = new JcePublicKeyKeyEncryptionMethodGenerator(pgpPublicKey)
                    .setProvider(new BouncyCastleProvider())
                    .setSecureRandom(new SecureRandom());
            cPk.addMethod(d);
            byte[] bytes = bOut.toByteArray();
            cOut = cPk.open(encryptFile, bytes.length);
            cOut.write(bytes);
        } catch (Exception e) {
            log.error("[encryptFile] error:", e);
            throw new BusinessException("500", "加密异常", ExceptionUtils.getStackTrace(e));
        } finally {
            try {
                if (cOut != null) {
                    cOut.close();
                }
                encryptFile.close();
            } catch (Exception e) {
                log.error("[encryptFile.close] error:", e);
            }
        }
    }


    public static boolean verifySignature(InputStream file, byte[] b, InputStream publicKey) throws GeneralSecurityException, IOException, PGPException {
        PGPObjectFactory pgpFact = new PGPObjectFactory(b);
        PGPSignatureList p3;
        Object o = pgpFact.nextObject();
        if (o instanceof PGPCompressedData) {
            PGPCompressedData c1 = (PGPCompressedData)o;
            pgpFact = new PGPObjectFactory(c1.getDataStream());
            p3 = (PGPSignatureList)pgpFact.nextObject();
        }
        else {
            p3 = (PGPSignatureList)o;
        }
        PGPPublicKeyRingCollection  pgpPubRingCollection = new PGPPublicKeyRingCollection(PGPUtil.getDecoderStream(publicKey));
        PGPSignature sig = p3.get(0);
        PGPPublicKey key = pgpPubRingCollection.getPublicKey(sig.getKeyID());
        sig.init(new JcaPGPContentVerifierBuilderProvider().setProvider(new BouncyCastleProvider()), key);
        int ch;
        while ((ch = file.read()) >= 0) {
            sig.update((byte)ch);
        }
        file.close();
        if (sig.verify()) {
            System.out.println("signature verified.");
            return true;
        }
        else {
            System.out.println("signature verification failed.");
            return false;
        }
    }
    public static byte[] createSignature(InputStream file, InputStream keyIn, char[] pass, boolean armor) throws GeneralSecurityException, IOException, PGPException {
        PGPSecretKey pgpSecKey = readSecretKey(keyIn);
        PGPPrivateKey pgpPrivKey = pgpSecKey.extractPrivateKey(new JcePBESecretKeyDecryptorBuilder().setProvider(new BouncyCastleProvider()).build(pass));
        PGPSignatureGenerator sGen = new PGPSignatureGenerator(new JcaPGPContentSignerBuilder(pgpSecKey.getPublicKey().getAlgorithm(), HashAlgorithmTags.SHA1).setProvider(new BouncyCastleProvider()));
        sGen.init(PGPSignature.BINARY_DOCUMENT, pgpPrivKey);
        ByteArrayOutputStream byteOut = new ByteArrayOutputStream();
        ArmoredOutputStream aOut = new ArmoredOutputStream(byteOut);
        BCPGOutputStream bOut = new BCPGOutputStream(byteOut);
        int ch;
        while ((ch = file.read()) >= 0) {
            sGen.update((byte)ch);
        }
        aOut.endClearText();
        file.close();
        sGen.generate().encode(bOut);
        if (armor) {
            aOut.close();
        }
        return byteOut.toByteArray();
    }

    private static PGPSecretKey readSecretKey(InputStream input) throws IOException, PGPException {
        PGPSecretKeyRingCollection pgpSec = new PGPSecretKeyRingCollection(PGPUtil.getDecoderStream(input));
        Iterator keyRingIter = pgpSec.getKeyRings();
        while (keyRingIter.hasNext()) {
            PGPSecretKeyRing keyRing = (PGPSecretKeyRing)keyRingIter.next();
            Iterator keyIter = keyRing.getSecretKeys();
            while (keyIter.hasNext()) {
                PGPSecretKey key = (PGPSecretKey)keyIter.next();
                if (key.isSigningKey()) {
                    return key;
                }
            }
        }
        throw new IllegalArgumentException("Can't find signing key in key ring.");
    }

    private static PGPPublicKey readPublicKey(InputStream publicKey) throws IOException, PGPException {
        PGPPublicKeyRingCollection pgpPub = new PGPPublicKeyRingCollection(PGPUtil.getDecoderStream(publicKey));
        PGPPublicKey key = null;
        Iterator rIt = pgpPub.getKeyRings();
        while (key == null && rIt.hasNext()) {
            PGPPublicKeyRing kRing = (PGPPublicKeyRing) rIt.next();
            Iterator kIt = kRing.getPublicKeys();
            while (key == null && kIt.hasNext()) {
                PGPPublicKey k = (PGPPublicKey) kIt.next();
                if (k.isEncryptionKey()) {
                    key = k;
                }
            }
        }
        if (key == null) {
            throw new IllegalArgumentException("[PGPHelper.readPublicKey] 公钥文件中未找到加密密钥。");
        }
        return key;
    }
    private static PGPPrivateKey findSecretKey(InputStream secretKey, long keyID, char[] pass) throws IOException, PGPException {
        PGPSecretKeyRingCollection pgpSec = new PGPSecretKeyRingCollection(PGPUtil.getDecoderStream(secretKey));
        PGPSecretKey pgpSecKey = pgpSec.getSecretKey(keyID);
        if (pgpSecKey == null) {
            return null;
        }
        PBESecretKeyDecryptor a = new JcePBESecretKeyDecryptorBuilder(new JcaPGPDigestCalculatorProviderBuilder()
                .setProvider(PROVIDER_NAME).build())
                .setProvider(PROVIDER_NAME).build(pass);
        return pgpSecKey.extractPrivateKey(a);
    }
    private PGPHelper() {
    }
}